#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SessionUaFDetector v1.2
Kali Linux 会话资源释放后重利用漏洞检测工具
检测场景: 用户注销后会话信息未被清除
法律声明: 仅限授权安全测试使用
"""
import sys
import os
import time
import json
import asyncio
import threading
import re
import random
import string
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Dict, List, Optional, Any, Tuple
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
import warnings
import argparse
import hashlib
import aiohttp

warnings.filterwarnings('ignore')

# ==================== 核心依赖 ====================
# 添加依赖检查函数
def check_dependencies():
    """检查并提示安装缺失的依赖"""
    missing = []
    try:
        from PyQt5.QtWidgets import QApplication
    except ImportError:
        missing.append("PyQt5")
    
    try:
        from rich.console import Console
    except ImportError:
        missing.append("rich")
    
    try:
        import aiolimiter
    except ImportError:
        missing.append("aiolimiter")
    
    try:
        import yaml
    except ImportError:
        missing.append("PyYAML")
    
    try:
        from prometheus_client import Counter, Histogram, Gauge, start_http_server
    except ImportError:
        missing.append("prometheus_client")
    
    if missing:
        print(f"[!] 缺少依赖包: {', '.join(missing)}")
        print(f"[*] 请运行: pip install {' '.join(missing)}")
        return False
    return True

try:
    from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                                 QHBoxLayout, QTabWidget, QLabel, QLineEdit,
                                 QPushButton, QTextEdit, QProgressBar, QComboBox,
                                 QFileDialog, QMessageBox, QTableWidget, QTableWidgetItem,
                                 QHeaderView, QSplitter, QPlainTextEdit, QCheckBox,
                                 QFrame, QGroupBox, QSpinBox, QDoubleSpinBox, QTextBrowser)
    from PyQt5.QtCore import pyqtSignal, QObject, Qt, QThread
    from PyQt5.QtGui import QFont, QTextCharFormat, QColor, QSyntaxHighlighter
    GUI_AVAILABLE = True
except ImportError:
    print("[-] PyQt5 未安装，GUI 模式不可用")
    GUI_AVAILABLE = False

try:
    from rich.console import Console
    from rich.table import Table
    from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
    from rich.panel import Panel
    from rich.syntax import Syntax
    from rich.traceback import install as rich_traceback_install
    rich_traceback_install()
    CLI_AVAILABLE = True
except ImportError:
    print("[-] Rich 库未安装，CLI 体验将降级")
    CLI_AVAILABLE = False

try:
    import aiolimiter
    AIOLIMITER_AVAILABLE = True
except ImportError:
    print("[-] aiolimiter 未安装，将使用基础速率限制")
    AIOLIMITER_AVAILABLE = False

try:
    import yaml
    YAML_AVAILABLE = True
except ImportError:
    print("[-] PyYAML 未安装，YAML配置功能禁用")
    YAML_AVAILABLE = False

METRICS_AVAILABLE = False
try:
    from prometheus_client import Counter, Histogram, Gauge, start_http_server
    METRICS_AVAILABLE = True
    print("[+] Prometheus 指标功能已启用")
except ImportError:
    print("[-] prometheus_client 未安装，监控功能禁用")

# ==================== 1. 企业级配置管理 ====================
class ConfigManager:
    """支持 YAML/JSON 的层级配置管理器"""
    DEFAULT_CONFIG = {
        "scanner": {
            "timeout": 15,
            "rate_limit": 5,
            "max_threads": 20,
            "retry_count": 3,
            "user_agents": [
                "SessionUaFDetector/1.2 (Kali Linux; Security Testing)",
                "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"
            ]
        },
        "detection": {
            "ua_threshold": 3,  # 重用尝试次数
            "login_success_indicators": [
                r'"success":\s*true', 
                r'"status":\s*"logged_in"',
                r'<div[^>]*class=["\']user-info["\']',
                r'Welcome\s+\w+',
                r'<a[^>]*href=["\']/logout["\']'
            ],
            "logout_success_indicators": [
                r'"success":\s*true',
                r'"status":\s*"logged_out"',
                r'<h1[^>]*>Logout Success',
                r'You have been logged out',
                r'logged out successfully'
            ],
            "protected_resource_indicators": [
                r'<div[^>]*class=["\']dashboard["\']',
                r'<a[^>]*href=["\']/admin["\']',
                r'"user_profile"',
                r'"account_settings"',
                r'<div[^>]*class=["\']user-panel["\']'
            ],
            "session_token_patterns": [
                r'sessionid=([a-fA-F0-9]+)',
                r'access_token=([a-zA-Z0-9\-\._~\+\/]+)',
                r'token=([a-zA-Z0-9\-\._~\+\/]+)',
                r'auth_token=([a-zA-Z0-9\-\._~\+\/]+)',
                r'jwt=([a-zA-Z0-9\-\._~\+\/]+)'
            ],
            "error_indicators": [
                r'unauthorized',
                r'not authenticated',
                r'not authorized',
                r'access denied',
                r'login required',
                r'invalid session',
                r'expired session'
            ]
        },
        "security": {
            "require_auth": True,
            "auto_sanitize_logs": True,
            "allowed_targets": []
        },
        "plugins": {
            "kali_tools": {
                "sqlmap": {"enabled": True, "path": "/usr/bin/sqlmap"},
                "nmap": {"enabled": False, "path": "/usr/bin/nmap"},
                "burp": {"enabled": True, "proxy": "127.0.0.1:8080"}
            }
        },
        "metrics": {
            "enabled": True,
            "port": 9091  # Different port for UaF detector
        }
    }

    def __init__(self, config_path: str = "~/.sessionuafd/config.yaml"):
        self.config_path = os.path.expanduser(config_path)
        self.config = self._load_config()

    def _load_config(self) -> Dict[str, Any]:
        """从 YAML/JSON 加载配置"""
        if not os.path.exists(self.config_path):
            os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
            self._save_config(self.DEFAULT_CONFIG)
            return self.DEFAULT_CONFIG
        
        try:
            with open(self.config_path, 'r') as f:
                content = f.read().strip()
                if not content:
                    print(f"[!] 配置文件 {self.config_path} 为空，使用默认配置")
                    return self.DEFAULT_CONFIG
                
                if self.config_path.endswith('.yaml') or self.config_path.endswith('.yml'):
                    if not YAML_AVAILABLE:
                        print(f"[!] YAML 支持库 (PyYAML) 未安装，跳过 {self.config_path}")
                        return self.DEFAULT_CONFIG
                    try:
                        config = yaml.safe_load(content) or {}
                        # 合并默认配置，确保所有键都存在
                        return self._merge_config(self.DEFAULT_CONFIG, config)
                    except Exception as e:
                        print(f"[!] YAML 文件解析失败: {e}")
                        return self.DEFAULT_CONFIG
                else:
                    try:
                        config = json.loads(content)
                        return self._merge_config(self.DEFAULT_CONFIG, config)
                    except json.JSONDecodeError as e:
                        print(f"[!] JSON 文件解析失败: {e}")
                        return self.DEFAULT_CONFIG
        except FileNotFoundError:
            print(f"[!] 配置文件 {self.config_path} 不存在，创建默认配置")
            self._save_config(self.DEFAULT_CONFIG)
            return self.DEFAULT_CONFIG
        except Exception as e:
            print(f"[!] 配置加载失败: {e}，使用默认配置")
            return self.DEFAULT_CONFIG

    def _merge_config(self, default: Dict, user: Dict) -> Dict:
        """递归合并配置，用户配置覆盖默认配置"""
        merged = default.copy()
        for key, value in user.items():
            if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
                merged[key] = self._merge_config(merged[key], value)
            else:
                merged[key] = value
        return merged

    def _save_config(self, config: Dict[str, Any]):
        """保存配置到文件"""
        try:
            os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
            with open(self.config_path, 'w') as f:
                if (self.config_path.endswith('.yaml') or self.config_path.endswith('.yml')) and YAML_AVAILABLE:
                    yaml.dump(config, f, default_flow_style=False, indent=2)
                else:
                    json.dump(config, f, indent=2)
        except Exception as e:
            print(f"[!] 配置保存失败: {e}")

    def get(self, key: str, default: Any = None) -> Any:
        """点分式获取配置值，如 scanner.timeout"""
        keys = key.split('.')
        value = self.config
        try:
            for k in keys:
                if isinstance(value, dict):
                    value = value.get(k)
                    if value is None:
                        return default
                else:
                    return default
            return value
        except (KeyError, TypeError):
            return default

    def set(self, key: str, value: Any):
        """设置配置值并自动保存"""
        keys = key.split('.')
        config = self.config
        try:
            for k in keys[:-1]:
                if not isinstance(config, dict):
                    # 如果路径不是字典，创建新字典
                    config = {}
                config = config.setdefault(k, {})
            
            if isinstance(config, dict):
                config[keys[-1]] = value
                self._save_config(self.config)
            else:
                print(f"[!] 无法设置配置项 {key}: 路径不是字典结构")
        except Exception as e:
            print(f"[!] 设置配置失败: {e}")

# ==================== 2. 增强的漏洞对象 ====================
class Vulnerability:
    """标准化的漏洞对象 - 增强版"""
    # 风险等级定义
    RISK_LEVELS = {
        "CRITICAL": {
            "score": 10.0,
            "color": "#dc3545",
            "description": "严重风险，需立即修复"
        },
        "HIGH": {
            "score": 8.0,
            "color": "#fd7e14",
            "description": "高风险，需尽快修复"
        },
        "MEDIUM": {
            "score": 5.0,
            "color": "#ffc107",
            "description": "中等风险，建议修复"
        },
        "LOW": {
            "score": 3.0,
            "color": "#28a745",
            "description": "低风险，可择期修复"
        },
        "INFO": {
            "score": 0.0,
            "color": "#6c757d",
            "description": "信息性发现"
        }
    }
    # 漏洞危害描述
    IMPACT_DESCRIPTIONS = {
        "Session Use-After-Free": {
            "impact": "攻击者可以利用已注销的会话令牌，冒充合法用户执行操作",
            "consequences": [
                "未经授权访问用户账户",
                "执行敏感操作（如转账、修改密码）",
                "窃取用户数据",
                "破坏数据完整性"
            ],
            "cvss_vector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:N"
        },
        "Token Reuse After Revocation": {
            "impact": "认证令牌在撤销后仍可被使用，导致身份验证绕过",
            "consequences": [
                "令牌失效机制失效",
                "持久化未授权访问",
                "账户接管"
            ],
            "cvss_vector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:N"
        },
        "Session Fixation": {
            "impact": "会话ID在登录前后保持不变，可能导致会话劫持",
            "consequences": [
                "会话劫持攻击",
                "中间人攻击",
                "账户接管"
            ],
            "cvss_vector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:L/A:N"
        }
    }
    # 修复建议
    REMEDIATION_ADVICE = {
        "Session Use-After-Free": {
            "immediate": [
                "确保注销操作后，服务器端会话对象被完全销毁",
                "清除客户端Cookie（如 sessionid）",
                "实现会话令牌一次性使用机制（如 Nonce）"
            ],
            "short_term": [
                "审查所有注销端点的实现逻辑",
                "记录所有会话创建和销毁事件",
                "实施会话状态同步（如使用 Redis 集群）"
            ],
            "long_term": [
                "采用更安全的会话管理机制（如 JWT with JTI, 或短期会话 + 刷新令牌）",
                "定期进行安全审计",
                "加强安全意识培训"
            ],
            "code_example": """
# Python Flask 示例：安全注销
from flask import session, request, jsonify
from your_session_store import delete_session  # 假设你有一个后端会话存储

@app.route('/logout', methods=['POST'])
def logout():
    session_id = session.get('session_id')
    if session_id:
        # 1. 从后端存储中删除会话
        delete_session(session_id)
        # 2. 清除当前请求的会话数据
        session.clear()
        # 3. (可选) 向客户端发送指令清除 Cookie
    return jsonify({'message': 'Logged out successfully'})

# Node.js Express 示例
app.post('/logout', (req, res) => {
  // 1. 从会话存储中删除会话
  req.sessionStore.destroy(req.sessionID, (err) => {
    if (err) {
      return res.status(500).json({ error: 'Logout failed' });
    }
    // 2. 销毁当前会话
    req.session.destroy(() => {
      // 3. 清除客户端 Cookie
      res.clearCookie('connect.sid');
      res.json({ message: 'Logged out successfully' });
    });
  });
});
            """
        },
        "Token Reuse After Revocation": {
            "immediate": [
                "实现令牌撤销列表（Token Revocation List）",
                "为每个令牌添加唯一标识符（JTI）并验证其有效性",
                "设置较短的令牌有效期"
            ],
            "short_term": [
                "实现令牌刷新机制",
                "使用分布式缓存（如Redis）存储已撤销令牌",
                "添加实时令牌状态检查"
            ],
            "long_term": [
                "采用OAuth 2.0授权框架",
                "实现基于证书的身份验证",
                "部署Web应用防火墙(WAF)监控可疑活动"
            ],
            "code_example": """
# Python JWT 撤销示例
import jwt
from datetime import datetime, timedelta

# 撤销令牌
revoked_tokens = set()

def revoke_token(token_jti):
    revoked_tokens.add(token_jti)

def is_token_revoked(token_jti):
    return token_jti in revoked_tokens

# 验证令牌
def verify_token(token):
    try:
        payload = jwt.decode(token, 'secret', algorithms=['HS256'])
        jti = payload.get('jti')
        if is_token_revoked(jti):
            return None  # 令牌已撤销
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None
            """
        }
    }

    def __init__(self, vuln_type: str, param: str, payload: str,
                 severity: str, evidence: str, confidence: float = 0.8,
                 test_details: Optional[Dict[str, Any]] = None):
        self.type = vuln_type
        self.param = param
        self.payload = payload
        self.severity = severity
        self.evidence = evidence
        self.confidence = confidence
        self.timestamp = datetime.now()
        self.test_details = test_details or {}
        # 计算风险评分
        self.risk_score = self._calculate_risk_score()
        # 获取危害描述
        self.impact = self._get_impact_description()
        # 获取修复建议
        self.remediation = self._get_remediation_advice()

    def _calculate_risk_score(self) -> float:
        """计算综合风险评分 (0-10)"""
        base_score = self.RISK_LEVELS.get(self.severity, {}).get("score", 0.0)
        # 根据置信度调整评分
        adjusted_score = base_score * self.confidence
        return round(adjusted_score, 2)

    def _get_impact_description(self) -> Dict[str, Any]:
        """获取漏洞危害描述"""
        return self.IMPACT_DESCRIPTIONS.get(self.type, {
            "impact": "未知影响",
            "consequences": ["需要进一步分析"],
            "cvss_vector": "N/A"
        })

    def _get_remediation_advice(self) -> Dict[str, Any]:
        """获取修复建议"""
        advice = self.REMEDIATION_ADVICE.get(self.type, {
            "immediate": ["请咨询安全专家"],
            "short_term": ["进行安全评估"],
            "long_term": ["建立安全开发流程"],
            "code_example": "# 请参考 OWASP 指南"
        })
        # 确保所有键都存在
        advice.setdefault("immediate", [])
        advice.setdefault("short_term", [])
        advice.setdefault("long_term", [])
        advice.setdefault("code_example", "# 请参考 OWASP 指南")
        return advice

    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": self.type,
            "param": self.param,
            "payload": self.payload,
            "severity": self.severity,
            "evidence": self.evidence,
            "confidence": self.confidence,
            "risk_score": self.risk_score,
            "timestamp": self.timestamp.isoformat(),
            "impact": self.impact,
            "remediation": self.remediation,
            "test_details": self.test_details
        }

    def __hash__(self):
        """用于去重"""
        return hash((self.type, self.param, self.payload))

    def __eq__(self, other):
        """用于去重"""
        if not isinstance(other, Vulnerability):
            return False
        return (self.type, self.param, self.payload) == (other.type, other.param, other.payload)

# ==================== 3. 测试过程记录器 ====================
class TestProcess:
    """测试过程详细记录"""
    def __init__(self):
        self.steps = []
        self.start_time = None
        self.end_time = None
        self.requests_sent = 0
        self.responses_received = 0
        self.errors_encountered = 0

    def start(self):
        """开始测试"""
        self.start_time = datetime.now()
        self.add_step("测试开始", "初始化扫描引擎", "INFO")

    def add_step(self, title: str, description: str, level: str = "INFO",
                 details: Optional[Dict[str, Any]] = None):
        """添加测试步骤"""
        step = {
            "timestamp": datetime.now().isoformat(),
            "title": title,
            "description": description,
            "level": level,
            "details": details or {}
        }
        self.steps.append(step)

    def end(self):
        """结束测试"""
        self.end_time = datetime.now()
        duration = (self.end_time - self.start_time).total_seconds()
        self.add_step(
            "测试完成",
            f"总耗时: {duration:.2f}秒, 发送请求: {self.requests_sent}, 接收响应: {self.responses_received}",
            "SUCCESS"
        )

    def to_html(self) -> str:
        """生成 HTML 格式的测试过程"""
        html = """
        <div class="test-process">
            <h3>📋 测试过程详情</h3>
            <div class="timeline">
        """
        for step in self.steps:
            level_colors = {
                "INFO": "#17a2b8",
                "SUCCESS": "#28a745",
                "WARNING": "#ffc107",
                "ERROR": "#dc3545"
            }
            color = level_colors.get(step['level'], "#6c757d")
            html += f"""
            <div class="timeline-item" style="border-left: 3px solid {color};">
                <div class="timeline-time">{step['timestamp']}</div>
                <div class="timeline-title">{step['title']}</div>
                <div class="timeline-description">{step['description']}</div>
            """
            # 修复: 检查details是否存在且非空
            if step.get('details') and isinstance(step['details'], dict):
                html += "<div class='timeline-details'>"
                for key, value in step['details'].items():
                    # 处理可能的非字符串值
                    display_value = str(value) if not isinstance(value, (dict, list)) else json.dumps(value)[:100] + "..."
                    html += f"<div><strong>{key}:</strong> {display_value}</div>"
                html += "</div>"
            html += "</div>"
        html += """
            </div>
        </div>
        """
        return html

# ==================== 4. UaF 检测策略 ====================
class UaFDetectionStrategy(ABC):
    """UaF 检测策略基类"""
    @abstractmethod
    def analyze(self, session_info: Dict[str, Any], logout_response: Dict[str, Any],
                reuse_attempts: List[Dict[str, Any]], context: Dict[str, Any]) -> Optional[Vulnerability]:
        """分析会话信息、注销响应和重用尝试，返回漏洞对象或 None"""
        pass

class SessionTokenReuseStrategy(UaFDetectionStrategy):
    """会话令牌重用检测策略"""
    def analyze(self, session_info: Dict[str, Any], logout_response: Dict[str, Any],
                reuse_attempts: List[Dict[str, Any]], context: Dict[str, Any]) -> Optional[Vulnerability]:
        test_process = context.get('test_process')
        config = context.get('config')
        
        # 修复: 确保配置项存在且为列表
        login_success_indicators = config.get("detection.login_success_indicators", []) or []
        logout_success_indicators = config.get("detection.logout_success_indicators", []) or []
        protected_resource_indicators = config.get("detection.protected_resource_indicators", []) or []
        session_token_patterns = config.get("detection.session_token_patterns", []) or []
        error_indicators = config.get("detection.error_indicators", []) or []

        if test_process:
            test_process.add_step(
                "UaF 策略分析",
                "检查注销后会话令牌是否失效",
                "INFO"
            )

        # 检查注销是否成功（可选，但有助于确认流程）
        logout_success = False
        logout_text = logout_response.get('text', '') or ''
        for indicator in logout_success_indicators:
            if re.search(indicator, logout_text, re.IGNORECASE):
                logout_success = True
                break
        
        if not logout_success:
            if test_process:
                test_process.add_step(
                    "注销检查",
                    "注销响应可能未成功，但继续测试会话重用",
                    "WARNING",
                    {"response_status": logout_response.get('status_code'), "response_text": logout_text[:200]}
                )

        # 检查重用尝试
        session_token = session_info.get('token', '') or ''
        session_cookies = session_info.get('cookies', {}) or {}
        
        # 修复: 确保reuse_attempts是列表且不为空
        if not isinstance(reuse_attempts, list) or not reuse_attempts:
            if test_process:
                test_process.add_step(
                    "UaF 检测",
                    "没有重用尝试数据",
                    "ERROR"
                )
            return None
        
        for i, attempt in enumerate(reuse_attempts):
            # 确保attempt是字典
            if not isinstance(attempt, dict):
                continue
                
            # 检查是否成功访问了受保护资源
            attempt_text = attempt.get('text', '') or ''
            for indicator in protected_resource_indicators:
                if re.search(indicator, attempt_text, re.IGNORECASE):
                    # 避免误报：检查是否是重定向到登录页面
                    is_login_page = False
                    for login_indicator in login_success_indicators:
                        if re.search(login_indicator, attempt_text, re.IGNORECASE):
                            is_login_page = True
                            break
                    
                    if is_login_page:
                        continue
                    
                    # 检查是否包含错误信息
                    has_error = False
                    for error_indicator in error_indicators:
                        if re.search(error_indicator, attempt_text, re.IGNORECASE):
                            has_error = True
                            break
                    
                    if has_error:
                        continue
                    
                    if test_process:
                        test_process.add_step(
                            f"UaF 检测成功 - 尝试 {i+1}",
                            f"使用已注销的会话成功访问受保护资源",
                            "ERROR",
                            {
                                "response_status": attempt.get('status_code'),
                                "response_text": attempt_text[:200],
                                "matched_indicator": indicator
                            }
                        )
                    
                    # 置信度基于重试次数 - 重试次数越多置信度越高
                    confidence = min(0.7 + (i * 0.1), 0.95)
                    
                    return Vulnerability(
                        vuln_type="Session Use-After-Free",
                        param="Session Token",
                        payload=session_token[:20] + "..." if session_token else "unknown",
                        severity="HIGH",
                        evidence=f"注销后第 {i+1} 次尝试使用旧会话令牌成功访问了受保护资源",
                        confidence=confidence,
                        test_details={
                            "session_token": session_token,
                            "logout_response_status": logout_response.get('status_code'),
                            "reuse_attempt_status": attempt.get('status_code'),
                            "reuse_attempt_text_snippet": attempt_text[:200],
                            "matched_indicator": indicator,
                            "attempt_number": i+1
                        }
                    )
            
            # 检查是否返回了新的会话令牌（表示旧令牌可能被误用或系统有缺陷）
            for pattern in session_token_patterns:
                match = re.search(pattern, attempt_text)
                if match:
                    new_token = match.group(1)
                    if session_token and new_token != session_token:
                        if test_process:
                            test_process.add_step(
                                f"UaF 检测成功 - 尝试 {i+1}",
                                f"使用已注销的会话获取了新的会话令牌",
                                "ERROR",
                                {
                                    "response_status": attempt.get('status_code'),
                                    "response_text": attempt_text[:200],
                                    "new_token": new_token[:10] + "...",
                                    "old_token": session_token[:10] + "..."
                                }
                            )
                        
                        return Vulnerability(
                            vuln_type="Session Use-After-Free",
                            param="Session Token",
                            payload=session_token[:20] + "..." if session_token else "unknown",
                            severity="HIGH",
                            evidence=f"注销后第 {i+1} 次尝试使用旧会话令牌获取了新的会话令牌",
                            confidence=0.9,
                            test_details={
                                "session_token": session_token,
                                "logout_response_status": logout_response.get('status_code'),
                                "reuse_attempt_status": attempt.get('status_code'),
                                "new_token_found": new_token,
                                "reuse_attempt_text_snippet": attempt_text[:200],
                                "attempt_number": i+1
                            }
                        )
        
        # 额外检查：如果所有重用尝试都返回相同的状态码（例如200），可能存在UaF
        if reuse_attempts:
            status_codes = [attempt.get('status_code', 0) for attempt in reuse_attempts if isinstance(attempt, dict)]
            baseline_status = logout_response.get('status_code', 0)
            
            # 如果注销返回200，且重用也返回200，可能存在UaF
            if baseline_status == 200 and status_codes and all(sc == 200 for sc in status_codes):
                # 但需要进一步确认不是静态页面
                is_static = False
                for attempt in reuse_attempts:
                    if not isinstance(attempt, dict):
                        continue
                    text = attempt.get('text', '')
                    if isinstance(text, str) and len(text) > 1000 and '<html' in text.lower():
                        # 检查是否包含登录元素
                        for indicator in login_success_indicators:
                            if re.search(indicator, text, re.IGNORECASE):
                                is_static = True
                                break
                        if is_static:
                            break
                
                if not is_static:
                    if test_process:
                        test_process.add_step(
                            "UaF 潜在检测",
                            "注销和重用请求均返回200状态码，可能存在UaF",
                            "WARNING",
                            {
                                "logout_status": baseline_status,
                                "reuse_statuses": status_codes
                            }
                        )
                    
                    return Vulnerability(
                        vuln_type="Session Use-After-Free",
                        param="HTTP Status Code",
                        payload="Status 200 after logout",
                        severity="MEDIUM",
                        evidence="注销和重用请求均返回200状态码，可能存在会话未清除",
                        confidence=0.6,
                        test_details={
                            "logout_status": baseline_status,
                            "reuse_statuses": status_codes,
                            "analysis": "状态码分析表明会话可能未正确终止"
                        }
                    )

        if test_process:
            test_process.add_step(
                "UaF 检测",
                "未发现会话重用成功的直接证据",
                "SUCCESS"
            )
        return None

class SessionFixationStrategy(UaFDetectionStrategy):
    """会话固定检测策略"""
    def analyze(self, session_info: Dict[str, Any], logout_response: Dict[str, Any],
                reuse_attempts: List[Dict[str, Any]], context: Dict[str, Any]) -> Optional[Vulnerability]:
        test_process = context.get('test_process')
        
        if not session_info or not session_info.get('initial_token'):
            return None
            
        initial_token = session_info['initial_token']
        reused_token = session_info.get('reused_token')
        
        if not reused_token:
            return None
            
        # 检查会话ID在注销和重用后是否保持不变
        if initial_token == reused_token:
            for i, attempt in enumerate(reuse_attempts):
                # 确保attempt是字典
                if not isinstance(attempt, dict):
                    continue
                    
                # 检查是否成功访问资源
                if attempt.get('status_code', 0) == 200:
                    if test_process:
                        test_process.add_step(
                            f"会话固定检测 - 尝试 {i+1}",
                            f"会话ID在注销后保持不变且访问成功",
                            "ERROR",
                            {
                                "initial_token": initial_token[:10] + "...",
                                "reused_token": reused_token[:10] + "...",
                                "status_code": attempt.get('status_code')
                            }
                        )
                    
                    return Vulnerability(
                        vuln_type="Session Fixation",
                        param="Session ID",
                        payload=f"Initial: {initial_token[:10]}..., Reused: {reused_token[:10]}...",
                        severity="HIGH",
                        evidence=f"会话ID在注销后保持不变，重用尝试 {i+1} 成功访问受保护资源",
                        confidence=0.85,
                        test_details={
                            "initial_session_id": initial_token,
                            "reused_session_id": reused_token,
                            "attempt_number": i+1,
                            "status_code": attempt.get('status_code')
                        }
                    )
        
        return None

# ==================== 5. UaF 异步扫描引擎 ====================
class SessionUaFScanner:
    """会话 UaF 扫描引擎"""
    def __init__(self, config: ConfigManager):
        self.config = config
        self.session = None
        # 修复: 检查aiolimiter是否可用
        if AIOLIMITER_AVAILABLE:
            self.limiter = aiolimiter.AsyncLimiter(
                config.get("scanner.rate_limit", 5),
                1
            )
        else:
            self.limiter = None
        self.strategies = [
            SessionTokenReuseStrategy(),
            SessionFixationStrategy()
        ]
        self.vulnerabilities = []
        self.metrics_available_instance = METRICS_AVAILABLE
        self.test_process = TestProcess()
        # 指标
        self.req_counter = None
        self.vuln_counter = None
        self.duration_hist = None
        if self.metrics_available_instance:
            try:
                self.req_counter = Counter('sessionuaftest_requests_total', 'Total requests', ['status'])
                self.vuln_counter = Counter('sessionuaftest_vulnerabilities_total', 'Vulnerabilities', ['type'])
                self.duration_hist = Histogram('sessionuaftest_request_duration_seconds', 'Request duration')
            except Exception as e:
                print(f"[!] 初始化 Prometheus 指标失败: {e}")
                self.metrics_available_instance = False

    async def _send_request(self, session: aiohttp.ClientSession,
                           url: str, method: str, headers: dict,
                           data: dict, request_id: str) -> Dict[str, Any]:
        """发送异步请求"""
        try:
            self.test_process.requests_sent += 1
            if not headers:
                headers = {}
            
            # 设置 User-Agent
            user_agents = self.config.get("scanner.user_agents", ["SessionUaFDetector/1.2 (Kali Linux; Security Testing)"])
            headers.setdefault("User-Agent", random.choice(user_agents))
            
            # 处理 Content-Type
            if isinstance(data, dict) and "Content-Type" not in headers:
                headers["Content-Type"] = "application/json"
            
            # 修复: 检查limiter是否可用
            if self.limiter:
                async with self.limiter:
                    result = await self._perform_request(session, url, method, headers, data)
            else:
                result = await self._perform_request(session, url, method, headers, data)
            
            # 记录指标
            if self.metrics_available_instance:
                try:
                    if self.req_counter:
                        self.req_counter.labels(status=result.get('status_code', 'error')).inc()
                    if self.duration_hist and 'duration' in result:
                        self.duration_hist.observe(result['duration'])
                except Exception as e:
                    print(f"[!] 记录 Prometheus 指标失败: {e}")
            
            return result
        except asyncio.TimeoutError:
            self.test_process.errors_encountered += 1
            self.test_process.add_step(
                f"请求超时 {request_id}",
                f"请求超时: {url}",
                "ERROR"
            )
            return {"error": "Request timeout"}
        except Exception as e:
            self.test_process.errors_encountered += 1
            self.test_process.add_step(
                f"请求失败 {request_id}",
                f"错误: {str(e)}",
                "ERROR"
            )
            return {"error": str(e)}
    
    async def _perform_request(self, session: aiohttp.ClientSession,
                              url: str, method: str, headers: dict,
                              data: dict) -> Dict[str, Any]:
        """执行实际请求"""
        start_time = time.time()
        if method == "GET":
            async with session.get(url, headers=headers, params=data,
                                  timeout=aiohttp.ClientTimeout(total=self.config.get("scanner.timeout", 15))) as resp:
                text = await resp.text()
                resp_headers = dict(resp.headers)
                status_code = resp.status
        else:
            if isinstance(data, dict):
                payload = json.dumps(data)
            else:
                payload = data
            
            async with session.request(method, url, headers=headers, data=payload,
                                     timeout=aiohttp.ClientTimeout(total=self.config.get("scanner.timeout", 15))) as resp:
                text = await resp.text()
                resp_headers = dict(resp.headers)
                status_code = resp.status
        
        end_time = time.time()
        duration = round(end_time - start_time, 3)
        self.test_process.responses_received += 1
        
        # 截断文本用于日志
        log_text = text[:200] if len(text) > 200 else text
        self.test_process.add_step(
            f"请求完成",
            f"状态码: {status_code}, 耗时: {duration}s, 长度: {len(text)}",
            "INFO",
            {
                "方法": method,
                "URL": url,
                "状态码": status_code,
                "响应长度": len(text),
                "耗时": f"{duration}s",
                "文本片段": log_text
            }
        )
        
        return {
            "status_code": status_code,
            "length": len(text),
            "text": text,
            "duration": duration,
            "headers": resp_headers,
            "cookies": dict(session.cookie_jar)
        }

    def _extract_session_info(self, login_response: Dict[str, Any], protected_response: Dict[str, Any], 
                             session: aiohttp.ClientSession) -> Dict[str, Any]:
        """提取会话信息"""
        session_info = {
            "token": None,
            "cookies": {},
            "initial_token": None
        }
        
        # 1. 从响应体中提取令牌
        login_success_indicators = self.config.get("detection.login_success_indicators", []) or []
        session_token_patterns = self.config.get("detection.session_token_patterns", []) or []
        
        login_text = login_response.get('text', '') or ''
        protected_text = protected_response.get('text', '') or ''
        
        # 检查登录响应
        for indicator in login_success_indicators:
            if re.search(indicator, login_text, re.IGNORECASE):
                for pattern in session_token_patterns:
                    match = re.search(pattern, login_text)
                    if match:
                        session_info['token'] = match.group(1)
                        session_info['initial_token'] = match.group(1)
                        break
                break
        
        # 检查受保护资源响应
        if not session_info['token']:
            for pattern in session_token_patterns:
                match = re.search(pattern, protected_text)
                if match:
                    session_info['token'] = match.group(1)
                    session_info['initial_token'] = match.group(1)
                    break
        
        # 2. 从Cookie中提取会话信息
        try:
            cookies = session.cookie_jar.filter_cookies("")
            session_info['cookies'] = {name: cookie.value for name, cookie in cookies.items()}
        except Exception as e:
            print(f"[!] 提取Cookie失败: {e}")
            session_info['cookies'] = {}
        
        # 3. 从响应头中提取令牌
        if not session_info['token']:
            headers = login_response.get('headers', {}) or {}
            auth_header = headers.get('Authorization', '')
            if auth_header and auth_header.startswith('Bearer '):
                session_info['token'] = auth_header.split('Bearer ')[1]
                session_info['initial_token'] = session_info['token']
        
        # 4. 提取常见会话ID
        if not session_info['token']:
            for name, value in session_info['cookies'].items():
                if 'session' in name.lower() or 'token' in name.lower():
                    session_info['token'] = value
                    session_info['initial_token'] = value
                    break
        
        return session_info

    async def scan(self, target: Dict[str, Any],
                  callback: Optional[callable] = None) -> List[Vulnerability]:
        """执行 UaF 扫描"""
        self.test_process.start()
        self.vulnerabilities = []
        
        # 验证必要参数
        required_fields = ['login_url', 'logout_url', 'protected_url']
        for field in required_fields:
            if field not in target or not target[field]:
                if callback: callback(f"错误: 缺少必要参数 '{field}'")
                self.test_process.add_step(
                    "参数验证",
                    f"缺少必要参数 '{field}'",
                    "ERROR"
                )
                self.test_process.end()
                return []
        
        # 设置请求方法
        login_method = target.get('login_method', 'POST').upper()
        logout_method = target.get('logout_method', 'POST').upper()
        protected_method = target.get('protected_method', 'GET').upper()
        
        # 合并headers
        base_headers = target.get('headers', {}) or {}
        login_headers = {**base_headers, **target.get('login_headers', {})}
        logout_headers = {**base_headers, **target.get('logout_headers', {})}
        protected_headers = {**base_headers, **target.get('protected_headers', {})}
        
        login_data = target.get('login_data', {})
        logout_data = target.get('logout_data', {})
        protected_data = target.get('protected_data', {})
        
        if callback: callback("开始执行会话UaF检测流程...")
        
        try:
            connector = aiohttp.TCPConnector(
                limit=self.config.get("scanner.max_threads", 20),
                ssl=False
            )
            timeout = aiohttp.ClientTimeout(total=self.config.get("scanner.timeout", 15))
            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
                # 1. 登录
                if callback: callback("步骤 1/4: 发送登录请求...")
                self.test_process.add_step(
                    "登录请求",
                    f"发送登录请求到 {target['login_url']}",
                    "INFO",
                    {
                        "method": login_method,
                        "url": target['login_url'],
                        "headers": {k: v for k, v in login_headers.items() if k.lower() not in ['authorization', 'cookie']},
                        "data": login_data
                    }
                )
                login_response = await self._send_request(session, target['login_url'], login_method, 
                                                         login_headers, login_data, "Login")
                if 'error' in login_response:
                    if callback: callback(f"登录请求失败: {login_response['error']}")
                    self.test_process.add_step(
                        "登录请求失败",
                        login_response['error'],
                        "ERROR"
                    )
                    self.test_process.end()
                    return []
                
                # 2. 验证登录成功
                login_success = False
                login_success_indicators = self.config.get("detection.login_success_indicators", []) or []
                login_text = login_response.get('text', '') or ''
                
                for indicator in login_success_indicators:
                    if re.search(indicator, login_text, re.IGNORECASE):
                        login_success = True
                        break
                
                if not login_success:
                    if callback: callback("警告: 未检测到登录成功的明确指示，将继续测试...")
                    self.test_process.add_step(
                        "登录状态检查",
                        "未检测到登录成功的明确指示，将继续测试",
                        "WARNING",
                        {
                            "status_code": login_response.get('status_code'),
                            "response_text": login_text[:200]
                        }
                    )
                
                # 3. 提取会话信息
                if callback: callback("提取会话信息...")
                session_info = self._extract_session_info(login_response, {}, session)
                
                if not session_info.get('token') and not session_info.get('cookies'):
                    if callback: callback("警告: 未找到会话标识符，测试可能不准确")
                    self.test_process.add_step(
                        "会话信息提取",
                        "未找到明确的会话标识符，测试可能不准确",
                        "WARNING"
                    )
                else:
                    token_display = session_info['token'][:10] + "..." if session_info.get('token') else "N/A"
                    cookies_display = ", ".join([f"{name}={value[:5]}..." for name, value in list(session_info['cookies'].items())[:3]])
                    self.test_process.add_step(
                        "会话信息提取",
                        f"成功提取会话信息",
                        "INFO",
                        {
                            "token": token_display,
                            "cookies": cookies_display,
                            "cookie_count": len(session_info['cookies'])
                        }
                    )
                
                # 4. 访问受保护资源（验证登录状态）
                if callback: callback("步骤 2/4: 访问受保护资源验证登录状态...")
                self.test_process.add_step(
                    "受保护资源访问",
                    f"访问 {target['protected_url']} 以验证登录状态",
                    "INFO"
                )
                protected_response = await self._send_request(session, target['protected_url'], 
                                                             protected_method, protected_headers, 
                                                             protected_data, "Protected-Before-Logout")
                
                if 'error' in protected_response:
                    if callback: callback(f"访问受保护资源失败: {protected_response['error']}")
                    self.test_process.add_step(
                        "受保护资源访问失败",
                        protected_response['error'],
                        "ERROR"
                    )
                    self.test_process.end()
                    return []
                
                # 5. 再次提取会话信息（可能在受保护资源请求中更新）
                session_info = self._extract_session_info(login_response, protected_response, session)
                
                # 6. 验证是否成功访问受保护资源
                protected_success = False
                protected_resource_indicators = self.config.get("detection.protected_resource_indicators", []) or []
                protected_text = protected_response.get('text', '') or ''
                
                for indicator in protected_resource_indicators:
                    if re.search(indicator, protected_text, re.IGNORECASE):
                        protected_success = True
                        break
                
                if not protected_success:
                    # 检查是否是重定向到登录页面
                    status_code = protected_response.get('status_code', 0)
                    location = protected_response.get('headers', {}).get('Location', '')
                    
                    if status_code in [301, 302, 307, 308] and location and ('login' in location.lower() or 'signin' in location.lower()):
                        if callback: callback("错误: 未成功访问受保护资源，可能登录失败或会话无效")
                        self.test_process.add_step(
                            "受保护资源访问",
                            "被重定向到登录页面，登录可能失败",
                            "ERROR",
                            {
                                "status_code": status_code,
                                "redirect_location": location
                            }
                        )
                        self.test_process.end()
                        return []
                    
                    if callback: callback("警告: 未检测到成功访问受保护资源的明确指示，将继续测试...")
                    self.test_process.add_step(
                        "受保护资源验证",
                        "未检测到成功访问的明确指示，将继续测试",
                        "WARNING",
                        {
                            "status_code": protected_response.get('status_code'),
                            "response_text": protected_text[:200]
                        }
                    )
                
                # 7. 注销
                if callback: callback("步骤 3/4: 发送注销请求...")
                self.test_process.add_step(
                    "注销请求",
                    f"发送注销请求到 {target['logout_url']}",
                    "INFO",
                    {
                        "method": logout_method,
                        "url": target['logout_url'],
                        "headers": {k: v for k, v in logout_headers.items() if k.lower() not in ['authorization', 'cookie']},
                        "data": logout_data
                    }
                )
                logout_response = await self._send_request(session, target['logout_url'], logout_method, 
                                                         logout_headers, logout_data, "Logout")
                
                if 'error' in logout_response:
                    if callback: callback(f"注销请求失败: {logout_response['error']}")
                    self.test_process.add_step(
                        "注销请求失败",
                        logout_response['error'],
                        "ERROR"
                    )
                    # 即使注销失败，也继续测试重用
                
                # 8. 尝试重用会话
                if callback: callback("步骤 4/4: 尝试重用已注销的会话...")
                self.test_process.add_step(
                    "会话重用测试",
                    "尝试重用已注销的会话访问受保护资源",
                    "INFO"
                )
                
                # 保存注销后的会话状态
                logout_session_info = self._extract_session_info({}, {}, session)
                session_info['reused_token'] = logout_session_info.get('token')
                
                # 执行多次重用尝试
                reuse_attempts = []
                ua_threshold = self.config.get("detection.ua_threshold", 3)
                
                for i in range(ua_threshold):
                    if callback: callback(f"重用尝试 {i+1}/{ua_threshold}...")
                    self.test_process.add_step(
                        f"会话重用尝试 {i+1}",
                        f"使用已注销的会话访问 {target['protected_url']}",
                        "INFO"
                    )
                    
                    # 重要：在重用尝试中，我们保持原始会话的cookie和headers
                    reuse_response = await self._send_request(session, target['protected_url'], 
                                                            protected_method, protected_headers, 
                                                            protected_data, f"Reuse-{i+1}")
                    reuse_attempts.append(reuse_response)
                    
                    # 短暂延迟，避免被速率限制
                    if i < ua_threshold - 1:
                        await asyncio.sleep(1)
                
                # 9. 策略分析
                self.test_process.add_step(
                    "UaF 分析",
                    "开始执行检测策略分析",
                    "INFO"
                )
                context = {
                    "config": self.config,
                    "test_process": self.test_process
                }
                potential_vulns = []
                
                for strategy in self.strategies:
                    strategy_name = strategy.__class__.__name__
                    self.test_process.add_step(
                        f"执行策略: {strategy_name}",
                        f"运行 {strategy_name} 检测",
                        "INFO"
                    )
                    vuln = strategy.analyze(session_info, logout_response, reuse_attempts, context)
                    if vuln:
                        potential_vulns.append(vuln)
                
                # 10. 去重和聚合
                self.vulnerabilities = self._deduplicate_and_aggregate(potential_vulns)
                
                # 11. 记录漏洞
                for vuln in self.vulnerabilities:
                    if self.metrics_available_instance:
                        try:
                            if self.vuln_counter:
                                self.vuln_counter.labels(type=vuln.type).inc()
                        except Exception as e:
                            print(f"[!] 记录 Prometheus 漏洞指标失败: {e}")
                    self.test_process.add_step(
                        f"发现漏洞: {vuln.type}",
                        vuln.evidence,
                        "ERROR",
                        {
                            "风险等级": vuln.severity,
                            "置信度": f"{vuln.confidence:.0%}",
                            "风险评分": vuln.risk_score
                        }
                    )
                    if callback: callback(f"发现漏洞: {vuln.type} (置信度: {vuln.confidence:.0%})")
                
                if not self.vulnerabilities:
                    if callback: callback("未发现会话UaF漏洞")
                    self.test_process.add_step(
                        "检测结果",
                        "未发现会话UaF漏洞",
                        "SUCCESS"
                    )
                
                self.test_process.end()
                
                if callback: 
                    vuln_count = len(self.vulnerabilities)
                    callback(f"UaF 扫描完成，发现 {vuln_count} 个漏洞")
                
                return self.vulnerabilities

        except Exception as e:
            print(f"[!] UaF 扫描过程中发生异常: {e}")
            import traceback
            traceback.print_exc()
            self.test_process.add_step(
                "UaF 扫描异常",
                str(e),
                "ERROR"
            )
            self.test_process.end()
            return []

    def _deduplicate_and_aggregate(self, vulns: List[Vulnerability]) -> List[Vulnerability]:
        """漏洞去重和聚合"""
        vuln_dict = {}
        for vuln in vulns:
            key = (vuln.type, vuln.param)
            if key not in vuln_dict or vuln.confidence > vuln_dict[key].confidence:
                vuln_dict[key] = vuln
        
        result = list(vuln_dict.values())
        result.sort(key=lambda v: v.risk_score, reverse=True)
        return result

# ==================== 6. Kali Linux 工具集成 ====================
class KaliToolIntegration:
    """Kali Linux 原生工具集成"""
    def __init__(self, config: ConfigManager):
        self.config = config
        self.tools = {
            "sqlmap": {
                "path": config.get("plugins.kali_tools.sqlmap.path", "/usr/bin/sqlmap"),
                "enabled": config.get("plugins.kali_tools.sqlmap.enabled", True),
                "category": "database"
            },
            "nmap": {
                "path": config.get("plugins.kali_tools.nmap.path", "/usr/bin/nmap"),
                "enabled": config.get("plugins.kali_tools.nmap.enabled", False),
                "category": "network"
            },
            "burp": {
                "proxy": config.get("plugins.kali_tools.burp.proxy", "127.0.0.1:8080"),
                "enabled": config.get("plugins.kali_tools.burp.enabled", True),
                "category": "proxy"
            }
        }
    
    def is_tool_available(self, tool_name: str) -> bool:
        """检查工具是否可用"""
        tool = self.tools.get(tool_name)
        if not tool or not tool['enabled']:
            return False
        
        if tool['category'] == 'proxy':
            import socket
            try:
                host, port = tool['proxy'].split(':')
                with socket.create_connection((host, int(port)), timeout=2):
                    pass
                return True
            except (ValueError, socket.gaierror, socket.timeout, OSError, ConnectionRefusedError):
                return False
        else:
            path = tool.get('path', '')
            return os.path.isfile(path) and os.access(path, os.X_OK)
    
    def import_from_burp(self, xml_file: str) -> List[Dict[str, Any]]:
        """从 Burp Suite XML 导入请求"""
        try:
            import xml.etree.ElementTree as ET
            tree = ET.parse(xml_file)
            root = tree.getroot()
            requests = []
            for item in root.findall('.//item'):
                url_elem = item.find('url')
                method_elem = item.find('method')
                if url_elem is None or method_elem is None:
                    print(f"[!] 跳过无效的 Burp 项: 缺少 URL 或 Method")
                    continue
                
                url = url_elem.text
                method = method_elem.text
                
                # 提取请求参数
                request_elem = item.find('request')
                if request_elem is not None and request_elem.text:
                    # 简单解析请求
                    request_text = request_elem.text
                    headers = {}
                    body = ""
                    
                    # 分割请求行和头部
                    parts = request_text.split('\r\n\r\n', 1)
                    if len(parts) > 1:
                        body = parts[1]
                        header_part = parts[0]
                        header_lines = header_part.split('\r\n')
                        for line in header_lines[1:]:  # 跳过请求行
                            if ':' in line:
                                key, value = line.split(':', 1)
                                headers[key.strip()] = value.strip()
                    
                    requests.append({
                        "url": url,
                        "method": method,
                        "headers": headers,
                        "body": body
                    })
            
            return requests
        except Exception as e:
            return [{"error": f"Burp 文件解析失败: {e}"}]

# ==================== 7. 企业级 GUI (Dark Mode) ====================
class DarkTheme:
    """Kali Linux 风格 Dark Mode 主题"""
    STYLESHEET = """
    QMainWindow {
        background-color: #1e1e1e;
        color: #d4d4d4;
    }
    QWidget {
        background-color: #1e1e1e;
        color: #d4d4d4;
    }
    QTabWidget::pane {
        border: 1px solid #3e3e42;
        background: #252526;
    }
    QTabBar::tab {
        background: #2d2d30;
        color: #cccccc;
        padding: 8px 20px;
        margin-right: 2px;
    }
    QTabBar::tab:selected {
        background: #007acc;
        color: white;
    }
    QPushButton {
        background-color: #007acc;
        color: white;
        padding: 8px 16px;
        border: none;
        border-radius: 4px;
        font-weight: bold;
    }
    QPushButton:hover {
        background-color: #005a9e;
    }
    QPushButton:disabled {
        background-color: #3e3e42;
        color: #6e6e70;
    }
    QLineEdit, QPlainTextEdit, QTextEdit {
        background-color: #2d2d30;
        color: #d4d4d4;
        border: 1px solid #3e3e42;
        padding: 6px;
        border-radius: 3px;
        selection-background-color: #264f78;
    }
    QTableWidget {
        background-color: #1e1e1e;
        color: #d4d4d4;
        gridline-color: #3e3e42;
        border: 1px solid #3e3e42;
    }
    QHeaderView::section {
        background-color: #2d2d30;
        color: #d4d4d4;
        padding: 6px;
        border: none;
        font-weight: bold;
    }
    QProgressBar {
        background-color: #3e3e42;
        border: none;
        height: 12px;
        border-radius: 6px;
    }
    QProgressBar::chunk {
        background-color: #007acc;
        border-radius: 6px;
    }
    QCheckBox {
        color: #d4d4d4;
        spacing: 8px;
    }
    QCheckBox::indicator {
        width: 18px;
        height: 18px;
        border: 1px solid #3e3e42;
        border-radius: 3px;
    }
    QCheckBox::indicator:checked {
        background-color: #007acc;
    }
    QLabel {
        color: #d4d4d4;
    }
    QGroupBox {
        border: 1px solid #3e3e42;
        border-radius: 4px;
        margin-top: 8px;
        padding-top: 8px;
    }
    QGroupBox::title {
        subcontrol-origin: margin;
        left: 10px;
        padding: 0 5px;
        color: #007acc;
        font-weight: bold;
    }
    QTextBrowser {
        background-color: #2d2d30;
        color: #d4d4d4;
        border: 1px solid #3e3e42;
    }
    """

class MainWindow(QMainWindow):
    """主窗口"""
    def __init__(self, config: ConfigManager):
        super().__init__()
        self.config = config
        self.scanner = None
        self.scan_thread = None
        self.scan_worker = None
        self.scan_is_running = False
        self.kali_tools = KaliToolIntegration(config)
        
        self.setWindowTitle("SessionUaFDetector v1.2 - Kali Linux Edition")
        self.setGeometry(100, 100, 1200, 800)
        self.setStyleSheet(DarkTheme.STYLESHEET)
        
        self._init_ui()
        
        # 启动指标服务器
        if METRICS_AVAILABLE and config.get("metrics.enabled", False):
            try:
                port = config.get("metrics.port", 9091)
                start_http_server(port)
                self.statusBar().showMessage(f"指标服务器运行在 :{port}")
            except Exception as e:
                print(f"[!] 启动 Prometheus 指标服务器失败: {e}")
                self.statusBar().showMessage("指标服务器启动失败")

    def _init_ui(self):
        """初始化UI"""
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        
        # 标签页
        self.tab_widget = QTabWidget()
        
        # 配置标签页
        config_tab = self._create_config_tab()
        self.tab_widget.addTab(config_tab, "⚙️ 配置")
        
        # 扫描标签页
        scan_tab = self._create_scan_tab()
        self.tab_widget.addTab(scan_tab, "🔍 扫描")
        
        # 结果标签页
        result_tab = self._create_result_tab()
        self.tab_widget.addTab(result_tab, "📊 结果")
        
        # 测试过程标签页
        self.process_tab = QTextBrowser()
        self.process_tab.setOpenExternalLinks(True)
        self.tab_widget.addTab(self.process_tab, "📋 测试过程")
        
        main_layout.addWidget(self.tab_widget)
        
        # 状态栏
        self.statusBar().showMessage("就绪")

    def _create_config_tab(self):
        """创建配置标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 目标配置
        target_group = QGroupBox("目标配置")
        target_layout = QVBoxLayout()
        
        # 登录URL
        target_layout.addWidget(QLabel("登录 URL:"))
        self.login_url_input = QLineEdit()
        self.login_url_input.setPlaceholderText("https://example.com/login")
        target_layout.addWidget(self.login_url_input)
        
        # 登录方法
        method_layout = QHBoxLayout()
        method_layout.addWidget(QLabel("登录方法:"))
        self.login_method_combo = QComboBox()
        self.login_method_combo.addItems(["POST", "GET"])
        method_layout.addWidget(self.login_method_combo)
        target_layout.addLayout(method_layout)
        
        # 登录数据
        target_layout.addWidget(QLabel("登录数据 (JSON):"))
        self.login_data_edit = QPlainTextEdit()
        self.login_data_edit.setPlainText(json.dumps({
            "username": "testuser",
            "password": "testpass"
        }, indent=2))
        self.login_data_edit.setMaximumHeight(100)
        target_layout.addWidget(self.login_data_edit)
        
        # 注销URL
        target_layout.addWidget(QLabel("注销 URL:"))
        self.logout_url_input = QLineEdit()
        self.logout_url_input.setPlaceholderText("https://example.com/logout")
        target_layout.addWidget(self.logout_url_input)
        
        # 注销方法
        logout_method_layout = QHBoxLayout()
        logout_method_layout.addWidget(QLabel("注销方法:"))
        self.logout_method_combo = QComboBox()
        self.logout_method_combo.addItems(["POST", "GET"])
        logout_method_layout.addWidget(self.logout_method_combo)
        target_layout.addLayout(logout_method_layout)
        
        # 受保护资源URL
        target_layout.addWidget(QLabel("受保护资源 URL:"))
        self.protected_url_input = QLineEdit()
        self.protected_url_input.setPlaceholderText("https://example.com/dashboard")
        target_layout.addWidget(self.protected_url_input)
        
        target_group.setLayout(target_layout)
        layout.addWidget(target_group)
        
        # 高级配置
        advanced_group = QGroupBox("高级配置")
        advanced_layout = QVBoxLayout()
        
        # 请求头
        advanced_layout.addWidget(QLabel("通用请求头 (JSON):"))
        self.headers_edit = QPlainTextEdit()
        self.headers_edit.setPlainText(json.dumps({
            "Content-Type": "application/json",
            "Accept": "application/json"
        }, indent=2))
        self.headers_edit.setMaximumHeight(100)
        advanced_layout.addWidget(self.headers_edit)
        
        # Kali工具集成
        advanced_layout.addWidget(QLabel("Kali 工具集成:"))
        kali_layout = QHBoxLayout()
        self.burp_check = QCheckBox("Burp Suite 代理")
        self.burp_check.setChecked(self.kali_tools.is_tool_available("burp"))
        kali_layout.addWidget(self.burp_check)
        
        import_btn = QPushButton("📥 导入 Burp 请求")
        import_btn.clicked.connect(self.import_burp_request)
        kali_layout.addWidget(import_btn)
        advanced_layout.addLayout(kali_layout)
        
        advanced_group.setLayout(advanced_layout)
        layout.addWidget(advanced_group)
        
        # 按钮
        btn_layout = QHBoxLayout()
        load_btn = QPushButton("📂 加载配置")
        load_btn.clicked.connect(self.load_config)
        save_btn = QPushButton("💾 保存配置")
        save_btn.clicked.connect(self.save_config)
        btn_layout.addWidget(load_btn)
        btn_layout.addWidget(save_btn)
        layout.addLayout(btn_layout)
        
        layout.addStretch()
        return tab

    def _create_scan_tab(self):
        """创建扫描标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 控制按钮
        btn_layout = QHBoxLayout()
        self.start_btn = QPushButton("🚀 开始检测")
        self.start_btn.clicked.connect(self.start_scan)
        self.stop_btn = QPushButton("⏹️ 停止")
        self.stop_btn.clicked.connect(self.stop_scan)
        self.stop_btn.setEnabled(False)
        btn_layout.addWidget(self.start_btn)
        btn_layout.addWidget(self.stop_btn)
        layout.addLayout(btn_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        layout.addWidget(self.progress_bar)
        
        # 日志
        layout.addWidget(QLabel("实时日志:"))
        self.log_edit = QTextEdit()
        self.log_edit.setReadOnly(True)
        self.log_edit.setMaximumHeight(200)
        layout.addWidget(self.log_edit)
        
        return tab

    def _create_result_tab(self):
        """创建结果标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 漏洞表格
        self.vuln_table = QTableWidget()
        self.vuln_table.setColumnCount(5)
        self.vuln_table.setHorizontalHeaderLabels(["类型", "参数", "严重程度", "风险评分", "置信度"])
        self.vuln_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.vuln_table.itemDoubleClicked.connect(self.show_vulnerability_details)
        layout.addWidget(self.vuln_table)
        
        # 详细信息
        layout.addWidget(QLabel("漏洞详情:"))
        self.vuln_details = QTextBrowser()
        self.vuln_details.setOpenExternalLinks(True)
        layout.addWidget(self.vuln_details)
        
        # 报告按钮
        report_btn = QPushButton("📝 生成报告")
        report_btn.clicked.connect(self.generate_report)
        layout.addWidget(report_btn)
        
        return tab

    def start_scan(self):
        """开始扫描"""
        if self.scan_is_running:
            QMessageBox.warning(self, "警告", "扫描已在进行中")
            return
        
        # 验证输入
        login_url = self.login_url_input.text().strip()
        logout_url = self.logout_url_input.text().strip()
        protected_url = self.protected_url_input.text().strip()
        
        if not all([login_url, logout_url, protected_url]):
            QMessageBox.warning(self, "错误", "请填写所有URL字段")
            return
        
        try:
            login_data = json.loads(self.login_data_edit.toPlainText())
        except json.JSONDecodeError:
            QMessageBox.warning(self, "错误", "登录数据必须是有效的JSON")
            return
        
        try:
            headers = json.loads(self.headers_edit.toPlainText())
        except json.JSONDecodeError:
            QMessageBox.warning(self, "错误", "请求头必须是有效的JSON")
            return
        
        # 准备目标
        target = {
            "login_url": login_url,
            "login_method": self.login_method_combo.currentText(),
            "login_data": login_data,
            "logout_url": logout_url,
            "logout_method": self.logout_method_combo.currentText(),
            "protected_url": protected_url,
            "protected_method": "GET",
            "headers": headers
        }
        
        # 重置UI
        self.vuln_table.setRowCount(0)
        self.vuln_details.clear()
        self.log_edit.clear()
        self.process_tab.clear()
        self.progress_bar.setValue(0)
        
        # 启动扫描
        self.scanner = SessionUaFScanner(self.config)
        self.scan_thread = QThread()
        self.scan_worker = ScanWorker(self.scanner, target)
        self.scan_worker.moveToThread(self.scan_thread)
        
        # 连接信号
        self.scan_worker.log_signal.connect(self.append_log)
        self.scan_worker.progress_signal.connect(self.progress_bar.setValue)
        self.scan_worker.result_signal.connect(self.on_vulnerability_found)
        self.scan_worker.finished_signal.connect(self.on_scan_finished)
        self.scan_worker.error_signal.connect(self.on_scan_error)
        
        self.scan_thread.started.connect(self.scan_worker.run)
        
        # 更新UI状态
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.scan_is_running = True
        self.statusBar().showMessage("扫描进行中...")
        
        # 启动线程
        self.scan_thread.start()

    def stop_scan(self):
        """停止扫描"""
        if self.scan_worker and self.scan_is_running:
            self.scan_worker.stop()
            self.scan_is_running = False

    def append_log(self, message: str):
        """追加日志"""
        self.log_edit.append(message)
        self.log_edit.verticalScrollBar().setValue(self.log_edit.verticalScrollBar().maximum())

    def on_vulnerability_found(self, vuln: Vulnerability):
        """发现漏洞"""
        row = self.vuln_table.rowCount()
        self.vuln_table.insertRow(row)
        
        self.vuln_table.setItem(row, 0, QTableWidgetItem(vuln.type))
        self.vuln_table.setItem(row, 1, QTableWidgetItem(vuln.param))
        
        # 设置严重程度颜色
        severity_item = QTableWidgetItem(vuln.severity)
        color = Vulnerability.RISK_LEVELS.get(vuln.severity, {}).get("color", "#6c757d")
        severity_item.setForeground(QColor(color))
        self.vuln_table.setItem(row, 2, severity_item)
        
        self.vuln_table.setItem(row, 3, QTableWidgetItem(f"{vuln.risk_score}/10"))
        self.vuln_table.setItem(row, 4, QTableWidgetItem(f"{vuln.confidence:.0%}"))
        
        # 存储完整漏洞对象
        self.vuln_table.item(row, 0).setData(Qt.UserRole, vuln)

    def show_vulnerability_details(self, item):
        """显示漏洞详情"""
        row = item.row()
        vuln_item = self.vuln_table.item(row, 0)
        if vuln_item:
            vuln = vuln_item.data(Qt.UserRole)
            if vuln:
                html = self._generate_vulnerability_html(vuln)
                self.vuln_details.setHtml(html)

    def _generate_vulnerability_html(self, vuln: Vulnerability) -> str:
        """生成漏洞详情HTML"""
        risk_level = Vulnerability.RISK_LEVELS.get(vuln.severity, {})
        color = risk_level.get("color", "#6c757d")
        
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <style>
                body {{
                    font-family: 'Segoe UI', sans-serif;
                    background: #1e1e1e;
                    color: #d4d4d4;
                    padding: 20px;
                }}
                .header {{
                    border-left: 5px solid {color};
                    padding-left: 15px;
                    margin-bottom: 20px;
                }}
                .header h1 {{
                    color: {color};
                    margin: 0;
                }}
                .section {{
                    background: #252526;
                    padding: 15px;
                    margin: 15px 0;
                    border-radius: 4px;
                    border: 1px solid #3e3e42;
                }}
                .section h3 {{
                    color: #007acc;
                    margin-top: 0;
                }}
                .metric {{
                    display: inline-block;
                    margin: 10px 20px 10px 0;
                }}
                .metric-label {{
                    color: #888;
                    font-size: 12px;
                }}
                .metric-value {{
                    font-size: 24px;
                    font-weight: bold;
                    color: {color};
                }}
                .code {{
                    background: #1e1e1e;
                    padding: 10px;
                    border-radius: 4px;
                    font-family: 'Courier New', monospace;
                    overflow-x: auto;
                    border: 1px solid #3e3e42;
                }}
                ul {{
                    padding-left: 20px;
                }}
                li {{
                    margin: 8px 0;
                }}
                .cvss {{
                    background: #2d2d30;
                    padding: 8px;
                    border-radius: 4px;
                    font-family: 'Courier New', monospace;
                    color: #4ade80;
                }}
            </style>
        </head>
        <body>
            <div class="header">
                <h1>[{vuln.severity}] {vuln.type}</h1>
                <p>{risk_level.get('description', '')}</p>
            </div>
            <div class="section">
                <h3>📊 风险评估</h3>
                <div class="metric">
                    <div class="metric-label">风险评分</div>
                    <div class="metric-value">{vuln.risk_score}/10</div>
                </div>
                <div class="metric">
                    <div class="metric-label">置信度</div>
                    <div class="metric-value">{vuln.confidence:.0%}</div>
                </div>
                <div class="metric">
                    <div class="metric-label">严重程度</div>
                    <div class="metric-value">{vuln.severity}</div>
                </div>
            </div>
            <div class="section">
                <h3>🔍 漏洞详情</h3>
                <p><strong>参数:</strong> {vuln.param}</p>
                <p><strong>Payload:</strong> <code>{vuln.payload}</code></p>
                <p><strong>证据:</strong> {vuln.evidence}</p>
                <p><strong>发现时间:</strong> {vuln.timestamp.strftime('%Y-%m-%d %H:%M:%S')}</p>
            </div>
            <div class="section">
                <h3>⚠️ 漏洞危害</h3>
                <p><strong>影响描述:</strong></p>
                <p>{vuln.impact.get('impact', '未知影响')}</p>
                <p><strong>潜在后果:</strong></p>
                <ul>
                    {''.join(f'<li>{c}</li>' for c in vuln.impact.get('consequences', []))}
                </ul>
                <p><strong>CVSS 向量:</strong></p>
                <div class="cvss">{vuln.impact.get('cvss_vector', 'N/A')}</div>
            </div>
            <div class="section">
                <h3>🛠️ 修复建议</h3>
                <p><strong>立即修复 (紧急):</strong></p>
                <ul>
                    {''.join(f'<li>{r}</li>' for r in vuln.remediation.get('immediate', []))}
                </ul>
                <p><strong>短期修复 (1-2周):</strong></p>
                <ul>
                    {''.join(f'<li>{r}</li>' for r in vuln.remediation.get('short_term', []))}
                </ul>
                <p><strong>长期改进 (1-3月):</strong></p>
                <ul>
                    {''.join(f'<li>{r}</li>' for r in vuln.remediation.get('long_term', []))}
                </ul>
            </div>
            <div class="section">
                <h3>💻 代码示例</h3>
                <div class="code">
                    <pre>{vuln.remediation.get('code_example', '请参考安全编码指南')}</pre>
                </div>
            </div>
            <div class="section">
                <h3>🔬 测试详情</h3>
                <pre>{json.dumps(vuln.test_details, indent=2, default=str)}</pre>
            </div>
        </body>
        </html>
        """
        return html

    def on_scan_finished(self, vulns: List[Vulnerability]):
        """扫描完成"""
        self.scan_is_running = False
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        
        # 显示测试过程
        if self.scanner and self.scanner.test_process:
            test_process_html = self.scanner.test_process.to_html()
            self.process_tab.setHtml(test_process_html)
        
        self.statusBar().showMessage(f"扫描完成，发现 {len(vulns)} 个漏洞")
        QMessageBox.information(self, "完成", f"扫描完成！发现 {len(vulns)} 个漏洞")

    def on_scan_error(self, error: str):
        """扫描错误"""
        self.scan_is_running = False
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.statusBar().showMessage(f"错误: {error}")
        QMessageBox.critical(self, "错误", f"扫描过程中出错:\n{error}")

    def generate_report(self):
        """生成报告"""
        # 修复: 检查scanner和test_process是否存在
        if not self.scanner or not self.scanner.test_process or not self.scanner.vulnerabilities:
            QMessageBox.warning(self, "警告", "没有可报告的扫描结果")
            return
        
        filename, _ = QFileDialog.getSaveFileName(
            self, "保存报告", f"session_uaf_report_{int(time.time())}.html", "HTML Files (*.html)"
        )
        
        if not filename:
            return
        
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                f.write(self._generate_report_html())
            QMessageBox.information(self, "成功", f"报告已保存至:\n{filename}")
            self.statusBar().showMessage(f"报告已保存至: {filename}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存报告失败:\n{str(e)}")

    def _generate_report_html(self) -> str:
        """生成报告HTML"""
        # 修复: 检查scanner是否存在
        if not self.scanner:
            return "<html><body>扫描器未初始化</body></html>"
            
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        vulns = self.scanner.vulnerabilities
        test_process = self.scanner.test_process
        
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="UTF-8">
            <title>SessionUaFDetector v1.2 扫描报告</title>
            <style>
                body {{
                    font-family: 'Segoe UI', sans-serif;
                    background: #f5f5f5;
                    margin: 0;
                    padding: 20px;
                }}
                .container {{
                    max-width: 1200px;
                    margin: 0 auto;
                    background: white;
                    padding: 40px;
                    border-radius: 8px;
                    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
                }}
                h1 {{
                    color: #007acc;
                    border-bottom: 3px solid #007acc;
                    padding-bottom: 15px;
                }}
                .summary {{
                    display: flex;
                    justify-content: space-around;
                    margin: 30px 0;
                }}
                .metric {{
                    text-align: center;
                    padding: 20px;
                    background: #f8f9fa;
                    border-radius: 8px;
                    flex: 1;
                    margin: 0 10px;
                }}
                .metric-value {{
                    font-size: 36px;
                    font-weight: bold;
                    color: #007acc;
                }}
                .metric-label {{
                    color: #666;
                    margin-top: 10px;
                }}
                .vuln-card {{
                    background: #fff;
                    border-radius: 4px;
                    box-shadow: 0 1px 3px rgba(0,0,0,0.1);
                    margin: 20px 0;
                    padding: 20px;
                }}
                .vuln-high {{ border-left: 5px solid #dc3545; }}
                .vuln-medium {{ border-left: 5px solid #ffc107; }}
                .vuln-low {{ border-left: 5px solid #28a745; }}
                .vuln-info {{ border-left: 5px solid #6c757d; }}
                .vuln-header {{
                    display: flex;
                    justify-content: space-between;
                    align-items: center;
                    margin-bottom: 15px;
                }}
                .vuln-title {{
                    font-size: 20px;
                    font-weight: bold;
                    color: #333;
                }}
                .risk-badge {{
                    padding: 5px 15px;
                    border-radius: 20px;
                    color: white;
                    font-weight: bold;
                }}
                .risk-high {{ background: #dc3545; }}
                .risk-medium {{ background: #ffc107; color: #333; }}
                .risk-low {{ background: #28a745; }}
                .risk-info {{ background: #6c757d; }}
                .code {{
                    background: #1e1e1e;
                    color: #d4d4d4;
                    padding: 15px;
                    border-radius: 4px;
                    font-family: 'Courier New', monospace;
                    overflow-x: auto;
                    margin: 10px 0;
                }}
                ul {{
                    padding-left: 20px;
                }}
                li {{
                    margin: 8px 0;
                }}
                .timeline {{
                    background: #f8f9fa;
                    padding: 20px;
                    border-radius: 4px;
                    margin: 20px 0;
                }}
                .timeline-item {{
                    padding: 10px;
                    margin: 5px 0;
                    border-left: 3px solid #007acc;
                }}
                .level-ERROR {{ border-left-color: #dc3545; }}
                .level-WARNING {{ border-left-color: #ffc107; }}
                .level-INFO {{ border-left-color: #17a2b8; }}
                .level-SUCCESS {{ border-left-color: #28a745; }}
            </style>
        </head>
        <body>
            <div class="container">
                <h1>🔍 SessionUaFDetector v1.2 扫描报告</h1>
                <p><strong>生成时间:</strong> {timestamp}</p>
                
                <div class="summary">
                    <div class="metric">
                        <div class="metric-value">{len(vulns)}</div>
                        <div class="metric-label">发现漏洞</div>
                    </div>
                    <div class="metric">
                        <div class="metric-value">{sum(1 for v in vulns if v.severity == 'HIGH')}</div>
                        <div class="metric-label">高危漏洞</div>
                    </div>
                    <div class="metric">
                        <div class="metric-value">{test_process.requests_sent}</div>
                        <div class="metric-label">发送请求</div>
                    </div>
                </div>
                
                <h2>📋 测试过程</h2>
                <div class="timeline">
        """
        
        # 添加测试过程
        for step in test_process.steps:
            level_class = f"level-{step['level']}"
            html += f"""
                    <div class="timeline-item {level_class}">
                        <strong>{step['timestamp']}</strong> - {step['title']}: {step['description']}
                    </div>
            """
        
        html += """
                </div>
                
                <h2>🔬 漏洞详情</h2>
        """
        
        # 添加漏洞详情
        for vuln in vulns:
            severity_class = f"vuln-{vuln.severity.lower()}"
            risk_class = f"risk-{vuln.severity.lower()}"
            html += f"""
                <div class="vuln-card {severity_class}">
                    <div class="vuln-header">
                        <div class="vuln-title">{vuln.type}</div>
                        <div class="risk-badge {risk_class}">{vuln.severity}</div>
                    </div>
                    <p><strong>风险评分:</strong> <span style="color: {Vulnerability.RISK_LEVELS.get(vuln.severity, {}).get('color', '#000')}; font-weight: bold;">{vuln.risk_score}/10</span></p>
                    <p><strong>置信度:</strong> {vuln.confidence:.0%}</p>
                    <p><strong>参数:</strong> <code>{vuln.param}</code></p>
                    <p><strong>证据:</strong> {vuln.evidence}</p>
                    
                    <h3>⚠️ 漏洞危害</h3>
                    <p>{vuln.impact.get('impact', '未知影响')}</p>
                    
                    <h3>🛠️ 修复建议</h3>
                    <p><strong>立即修复:</strong></p>
                    <ul>
                        {''.join(f'<li>{r}</li>' for r in vuln.remediation.get('immediate', []))}
                    </ul>
                    
                    <h3>💻 代码示例</h3>
                    <div class="code">
                        <pre>{vuln.remediation.get('code_example', '请参考安全编码指南')}</pre>
                    </div>
                </div>
            """
        
        html += """
            </div>
        </body>
        </html>
        """
        return html

    def import_burp_request(self):
        """导入Burp请求"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择 Burp Suite XML 文件", "", "XML Files (*.xml)"
        )
        
        if not file_path:
            return
        
        requests = self.kali_tools.import_from_burp(file_path)
        
        # 修复: 检查requests是否为空或包含错误
        if not requests or (isinstance(requests, list) and len(requests) > 0 and 'error' in requests[0]):
            error_msg = requests[0].get('error', '未知错误') if requests else '未知错误'
            QMessageBox.warning(self, "错误", f"导入失败: {error_msg}")
            return
        
        # 假设第一个请求是登录请求
        if requests and len(requests) > 0:
            req = requests[0]
            self.login_url_input.setText(req['url'])
            self.login_method_combo.setCurrentText(req['method'])
            
            # 提取登录参数
            if 'body' in req and req['body']:
                try:
                    # 尝试解析JSON
                    params = json.loads(req['body'])
                    self.login_data_edit.setPlainText(json.dumps(params, indent=2))
                except:
                    # 尝试解析表单数据
                    if 'Content-Type' in req['headers'] and 'application/x-www-form-urlencoded' in req['headers']['Content-Type']:
                        params = parse_qs(req['body'])
                        params = {k: v[0] if isinstance(v, list) and v else v for k, v in params.items()}
                        self.login_data_edit.setPlainText(json.dumps(params, indent=2))
            
            QMessageBox.information(self, "成功", "请求导入成功！请检查并完善其他配置。")

    def load_config(self):
        """加载配置"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择配置文件", "", "YAML Files (*.yaml *.yml);;JSON Files (*.json)"
        )
        
        if not file_path:
            return
        
        try:
            with open(file_path, 'r') as f:
                content = f.read()
                if file_path.endswith('.yaml') or file_path.endswith('.yml'):
                    if not YAML_AVAILABLE:
                        QMessageBox.warning(self, "警告", "PyYAML未安装，无法加载YAML文件")
                        return
                    import yaml
                    config = yaml.safe_load(content)
                else:
                    config = json.loads(content)
            
            # 应用配置
            if 'target' in config:
                target = config['target']
                self.login_url_input.setText(target.get('login_url', ''))
                self.logout_url_input.setText(target.get('logout_url', ''))
                self.protected_url_input.setText(target.get('protected_url', ''))
                
                if 'login_data' in target:
                    self.login_data_edit.setPlainText(json.dumps(target['login_data'], indent=2))
                
                if 'headers' in target:
                    self.headers_edit.setPlainText(json.dumps(target['headers'], indent=2))
            
            QMessageBox.information(self, "成功", "配置加载成功！")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"加载配置失败:\n{str(e)}")

    def save_config(self):
        """保存配置"""
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存配置文件", "", "YAML Files (*.yaml *.yml);;JSON Files (*.json)"
        )
        
        if not file_path:
            return
        
        try:
            config = {
                "target": {
                    "login_url": self.login_url_input.text().strip(),
                    "login_method": self.login_method_combo.currentText(),
                    "logout_url": self.logout_url_input.text().strip(),
                    "logout_method": self.logout_method_combo.currentText(),
                    "protected_url": self.protected_url_input.text().strip(),
                    "protected_method": "GET"
                }
            }
            
            try:
                config["target"]["login_data"] = json.loads(self.login_data_edit.toPlainText())
                config["target"]["headers"] = json.loads(self.headers_edit.toPlainText())
            except json.JSONDecodeError:
                QMessageBox.warning(self, "警告", "配置包含无效的JSON，将保存为原始文本")
                config["target"]["login_data_raw"] = self.login_data_edit.toPlainText()
                config["target"]["headers_raw"] = self.headers_edit.toPlainText()
            
            with open(file_path, 'w') as f:
                if file_path.endswith('.yaml') or file_path.endswith('.yml'):
                    if not YAML_AVAILABLE:
                        QMessageBox.warning(self, "警告", "PyYAML未安装，将保存为JSON格式")
                        json.dump(config, f, indent=2)
                    else:
                        import yaml
                        yaml.dump(config, f, default_flow_style=False, indent=2)
                else:
                    json.dump(config, f, indent=2)
            
            QMessageBox.information(self, "成功", "配置保存成功！")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存配置失败:\n{str(e)}")

class ScanWorker(QObject):
    """扫描工作线程"""
    log_signal = pyqtSignal(str)
    progress_signal = pyqtSignal(int)
    result_signal = pyqtSignal(Vulnerability)
    finished_signal = pyqtSignal(list)
    error_signal = pyqtSignal(str)
    stop_signal = pyqtSignal()
    
    def __init__(self, scanner: SessionUaFScanner, target: Dict[str, Any]):
        super().__init__()
        self.scanner = scanner
        self.target = target
        self._stop_flag = False
        self.stop_signal.connect(self._set_stop_flag)
    
    def _set_stop_flag(self):
        self._stop_flag = True
    
    def stop(self):
        self.stop_signal.emit()
    
    def run(self):
        """执行扫描"""
        # 修复: 在QThread中正确使用asyncio
        loop = None
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            def progress_callback(msg: str):
                if self._stop_flag:
                    return False
                self.log_signal.emit(msg)
                return True
            
            self.log_signal.emit("初始化扫描引擎...")
            self.progress_signal.emit(10)
            
            if self._stop_flag:
                self.finished_signal.emit([])
                return
            
            self.log_signal.emit("开始执行UaF检测流程...")
            vulns = loop.run_until_complete(self.scanner.scan(self.target, progress_callback))
            
            if self._stop_flag:
                self.log_signal.emit("扫描已停止")
                self.finished_signal.emit([])
            else:
                self.progress_signal.emit(100)
                self.finished_signal.emit(vulns)
        
        except Exception as e:
            self.error_signal.emit(str(e))
        finally:
            if loop:
                loop.close()

# ==================== 8. CLI 界面 ====================
class CLIApp:
    """命令行界面"""
    def __init__(self, config: ConfigManager):
        self.config = config
        self.console = Console() if CLI_AVAILABLE else None
        self.scanner = SessionUaFScanner(config)
        self.kali_tools = KaliToolIntegration(config)
    
    def print_banner(self):
        """打印Banner"""
        if self.console:
            banner = r"""
[bold blue]    ____  ____________    __  __            __           _   _            
   / __ \/ ____/ ____/   / / / /_  ______  / /____  _____| | | | ___  ___  
  / /_/ / /   / __/     / /_/ / / / / __ \/ __/ _ \/ ___/ |_| |/ _ \/ __| 
 / _, _/ /___/ /___    / __  / /_/ / / / / /_/  __/ /  |  _  | (_) \__ \ 
/_/ |_|\____/_____/   /_/ /_/\__,_/_/ /_/\__/\___/_/   |_| |_|\___/|___/ 
[/bold blue]
[bold yellow]:: Kali Linux Edition :: v1.2 :: Session UaF Scanner ::[/bold yellow]
            """
            self.console.print(Panel(banner, style="bold white"))
        else:
            print(r"""
    ____  ____________    __  __            __           _   _            
   / __ \/ ____/ ____/   / / / /_  ______  / /____  _____| | | | ___  ___  
  / /_/ / /   / __/     / /_/ / / / / __ \/ __/ _ \/ ___/ |_| |/ _ \/ __| 
 / _, _/ /___/ /___    / __  / /_/ / / / / /_/  __/ /  |  _  | (_) \__ \ 
/_/ |_|\____/_____/   /_/ /_/\__,_/_/ /_/\__/\___/_/   |_| |_|\___/|___/ 
:: Kali Linux Edition :: v1.2 :: Session UaF Scanner ::
            """)
    
    def run_interactive(self):
        """交互式模式"""
        self.print_banner()
        
        # 法律确认
        if self.config.get("security.require_auth", True):
            if self.console:
                self.console.print("\n[bold red]⚠️  法律警告[/bold red]")
                self.console.print("本工具仅用于授权的安全测试")
                confirm = input("\n输入 'YES' 确认授权: ")
            else:
                print("\n⚠️  法律警告")
                print("本工具仅用于授权的安全测试")
                confirm = input("\n输入 'YES' 确认授权: ")
            
            if confirm.strip().upper() != "YES":
                if self.console:
                    self.console.print("[red]未授权，退出。[/red]")
                else:
                    print("未授权，退出。")
                return
        
        # 收集目标信息
        login_url = input("登录 URL: ").strip()
        logout_url = input("注销 URL: ").strip()
        protected_url = input("受保护资源 URL: ").strip()
        
        if not all([login_url, logout_url, protected_url]):
            print("错误: 所有URL字段均为必填")
            return
        
        login_method = input("登录 HTTP 方法 [POST]: ").strip().upper() or "POST"
        logout_method = input("注销 HTTP 方法 [POST]: ").strip().upper() or "POST"
        
        # 收集登录数据
        print("\n输入登录请求数据 (JSON 格式)，按 Ctrl+D (或 Ctrl+Z) 结束:")
        login_data_lines = []
        try:
            while True:
                line = input()
                login_data_lines.append(line)
        except EOFError:
            pass
        
        try:
            login_data = json.loads("\n".join(login_data_lines))
        except json.JSONDecodeError:
            print("错误: 无效的JSON格式")
            return
        
        # 执行扫描
        print("\n开始执行会话UaF检测...")
        target = {
            "login_url": login_url,
            "login_method": login_method,
            "login_data": login_data,
            "logout_url": logout_url,
            "logout_method": logout_method,
            "protected_url": protected_url,
            "protected_method": "GET"
        }
        
        if CLI_AVAILABLE and self.console:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TaskProgressColumn(),
                console=self.console
            ) as progress:
                task = progress.add_task("[cyan]正在扫描...", total=100)
                
                def progress_callback(msg: str):
                    progress.update(task, advance=25, description=msg)
                    return True
                
                vulns = asyncio.run(self.scanner.scan(target, progress_callback))
                progress.update(task, completed=100)
        else:
            def simple_progress(msg: str):
                print(f"  {msg}")
                return True
            
            print("正在扫描...")
            vulns = asyncio.run(self.scanner.scan(target, simple_progress))
            print("扫描完成!")
        
        # 显示结果
        self.display_results(vulns)
        
        # 保存报告
        if vulns:
            save = input("\n是否保存报告? (y/n): ").strip().lower()
            if save == 'y':
                report_file = self.save_report(vulns)
                print(f"报告已保存至: {report_file}")
    
    def display_results(self, vulns: List[Vulnerability]):
        """显示扫描结果"""
        if not vulns:
            if self.console:
                self.console.print("[green]✅ 未发现会话UaF漏洞[/green]")
            else:
                print("✅ 未发现会话UaF漏洞")
            return
        
        if self.console:
            table = Table(title="漏洞详情", show_header=True, header_style="bold magenta")
            table.add_column("类型", style="cyan")
            table.add_column("参数", style="yellow")
            table.add_column("严重程度", style="red")
            table.add_column("风险评分", justify="right")
            table.add_column("置信度", justify="right")
            table.add_column("证据", style="dim")
            
            for vuln in vulns:
                severity_color = {
                    "CRITICAL": "bold red",
                    "HIGH": "red",
                    "MEDIUM": "yellow",
                    "LOW": "green",
                    "INFO": "dim"
                }.get(vuln.severity, "white")
                
                table.add_row(
                    vuln.type,
                    vuln.param,
                    f"[{severity_color}]{vuln.severity}[/{severity_color}]",
                    f"{vuln.risk_score:.1f}/10",
                    f"{vuln.confidence:.0%}",
                    vuln.evidence[:60] + "..." if len(vuln.evidence) > 60 else vuln.evidence
                )
            
            self.console.print(table)
            self.console.print(f"\n[bold red]⚠️  发现 {len(vulns)} 个会话UaF漏洞![/bold red]")
        else:
            print(f"\n⚠️  发现 {len(vulns)} 个会话UaF漏洞!")
            for i, vuln in enumerate(vulns):
                print(f"  [{i+1}] [{vuln.severity}] {vuln.type}")
                print(f"      参数: {vuln.param}")
                print(f"      风险评分: {vuln.risk_score}/10")
                print(f"      置信度: {vuln.confidence:.0%}")
                print(f"      证据: {vuln.evidence}")
    
    def save_report(self, vulns: List[Vulnerability]) -> str:
        """保存报告"""
        filename = f"session_uaf_report_{int(time.time())}.html"
        
        # 生成HTML报告
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="UTF-8">
            <title>SessionUaFDetector CLI 报告</title>
            <style>
                body {{
                    font-family: Arial, sans-serif;
                    margin: 40px;
                    line-height: 1.6;
                }}
                .container {{
                    max-width: 1200px;
                    margin: 0 auto;
                }}
                h1 {{
                    color: #007acc;
                    border-bottom: 2px solid #007acc;
                    padding-bottom: 10px;
                }}
                .summary {{
                    background: #f5f5f5;
                    padding: 20px;
                    border-radius: 5px;
                    margin: 20px 0;
                }}
                .vuln {{
                    border: 1px solid #e1e1e1;
                    border-radius: 5px;
                    padding: 15px;
                    margin: 15px 0;
                }}
                .vuln-high {{
                    border-left: 5px solid #dc3545;
                }}
                .vuln-medium {{
                    border-left: 5px solid #ffc107;
                }}
                .vuln-low {{
                    border-left: 5px solid #28a745;
                }}
                .code {{
                    background: #f8f9fa;
                    padding: 10px;
                    border-radius: 4px;
                    font-family: monospace;
                    overflow-x: auto;
                }}
                ul {{
                    padding-left: 20px;
                }}
                li {{
                    margin: 5px 0;
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <h1>SessionUaFDetector v1.2 CLI 报告</h1>
                <p><strong>生成时间:</strong> {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
                
                <div class="summary">
                    <h2>摘要</h2>
                    <p><strong>发现漏洞:</strong> {len(vulns)}</p>
                    <p><strong>高危漏洞:</strong> {sum(1 for v in vulns if v.severity == 'HIGH')}</p>
                </div>
                
                <h2>漏洞详情</h2>
        """
        
        for vuln in vulns:
            severity_class = f"vuln-{vuln.severity.lower()}"
            html += f"""
                <div class="vuln {severity_class}">
                    <h3>[{vuln.severity}] {vuln.type}</h3>
                    <p><strong>参数:</strong> {vuln.param}</p>
                    <p><strong>风险评分:</strong> {vuln.risk_score}/10</p>
                    <p><strong>置信度:</strong> {vuln.confidence:.0%}</p>
                    <p><strong>证据:</strong> {vuln.evidence}</p>
                    
                    <h4>漏洞危害</h4>
                    <p>{vuln.impact.get('impact', '未知影响')}</p>
                    
                    <h4>修复建议</h4>
                    <ul>
                        {''.join(f'<li>{r}</li>' for r in vuln.remediation.get('immediate', []))}
                    </ul>
                    
                    <h4>代码示例</h4>
                    <div class="code">
                        <pre>{vuln.remediation.get('code_example', '请参考安全编码指南')}</pre>
                    </div>
                </div>
            """
        
        html += """
            </div>
        </body>
        </html>
        """
        
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                f.write(html)
            return filename
        except Exception as e:
            print(f"保存报告失败: {e}")
            return "保存失败"

# ==================== 9. 主入口 ====================
def main():
    """程序主入口"""
    # 检查依赖
    if not check_dependencies():
        response = input("是否继续运行？(部分功能将受限) [y/N]: ")
        if response.lower() != 'y':
            sys.exit(1)
    
    parser = argparse.ArgumentParser(
        description="SessionUaFDetector v1.2 - 会话资源释放后重利用漏洞检测工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # GUI 模式
  python3 session_uaf_detector.py --gui
  
  # CLI 交互模式
  python3 session_uaf_detector.py --cli
  
  # 快速扫描
  python3 session_uaf_detector.py --login-url https://example.com/login --logout-url https://example.com/logout --protected-url https://example.com/dashboard
  
  # 加载配置文件
  python3 session_uaf_detector.py --config ~/.sessionuafd/config.yaml --gui
        """
    )
    parser.add_argument("--gui", action="store_true", help="启动图形界面")
    parser.add_argument("--cli", action="store_true", help="启动命令行界面")
    parser.add_argument("--login-url", help="登录URL")
    parser.add_argument("--logout-url", help="注销URL")
    parser.add_argument("--protected-url", help="受保护资源URL")
    parser.add_argument("--config", help="配置文件路径")
    parser.add_argument("--auth", help="授权令牌")
    
    args = parser.parse_args()
    config = ConfigManager(args.config or "~/.sessionuafd/config.yaml")
    
    if args.auth:
        config.set("security.auth_token", args.auth)
    
    # 检查Kali Linux环境
    if os.path.exists("/etc/os-release"):
        try:
            with open("/etc/os-release") as f:
                if "kali" not in f.read().lower():
                    print("[!] 警告: 未在 Kali Linux 中运行，部分功能可能受限")
        except Exception as e:
            print(f"[!] 检查系统信息时出错: {e}")
    
    # 启动GUI
    if args.gui and GUI_AVAILABLE:
        app = QApplication(sys.argv)
        app.setStyle("Fusion")
        window = MainWindow(config)
        window.show()
        sys.exit(app.exec_())
    
    # 启动CLI
    elif args.cli or (len(sys.argv) == 1 and not args.gui):
        cli = CLIApp(config)
        cli.run_interactive()
    
    # 快速扫描
    elif args.login_url and args.logout_url and args.protected_url:
        cli = CLIApp(config)
        target = {
            "login_url": args.login_url,
            "logout_url": args.logout_url,
            "protected_url": args.protected_url,
            "login_method": "POST",
            "logout_method": "POST",
            "protected_method": "GET",
            "login_data": {"username": "test", "password": "test"}
        }
        vulns = asyncio.run(cli.scanner.scan(target))
        cli.display_results(vulns)
        
        if vulns:
            save = input("\n是否保存报告? (y/n): ").strip().lower()
            if save == 'y':
                report_file = cli.save_report(vulns)
                print(f"报告已保存至: {report_file}")
    
    else:
        parser.print_help()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n用户中断，退出。")
        sys.exit(0)
    except Exception as e:
        print(f"\n致命错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
